package com.zhm.demo.rocksdb;

import com.google.common.collect.Lists;
import lombok.AllArgsConstructor;
import lombok.Data;
import org.rocksdb.*;
import org.rocksdb.util.SizeUnit;

import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.stream.Collectors;

import static java.nio.charset.StandardCharsets.UTF_8;

public class RocksDBExample {
    private static final String dbPath = "./data-db/";
    private static final String cfdbPath = "./data-cf/";
    private static final String txdbPath = "./data-tx/";


    public static RocksDB db;
    public static final Map<String, ColumnFamilyHandle> columnFamilyHandles = new HashMap<>();

    public static Options getOptions() {
        final Options options = new Options();
        final Filter bloomFilter = new BloomFilter(10);
        final ReadOptions readOptions = new ReadOptions().setFillCache(false);
        final Statistics stats = new Statistics();
        final RateLimiter rateLimiter = new RateLimiter(10000000, 10000, 10);

        options.setCreateIfMissing(true)
                .setStatistics(stats)
                .setWriteBufferSize(8 * SizeUnit.KB)
                .setMaxWriteBufferNumber(3)
                .setMaxBackgroundJobs(10)
                .setCompressionType(CompressionType.SNAPPY_COMPRESSION)
                .setCompactionStyle(CompactionStyle.UNIVERSAL);

        final BlockBasedTableConfig table_options = new BlockBasedTableConfig();
        Cache cache = new LRUCache(64 * 1024, 6);
        table_options.setBlockCache(cache)
                .setFilterPolicy(bloomFilter)
                .setBlockSizeDeviation(5)
                .setBlockRestartInterval(10)
                .setCacheIndexAndFilterBlocks(true)
                .setBlockCacheCompressed(new LRUCache(64 * 1000, 10));
        options.setTableFormatConfig(table_options);
        options.setRateLimiter(rateLimiter);
        return options;
    }

    public static void main(String[] args) {
//        testBasicOperate();
//        testCustomColumnFamily();
//        testTransaction();
        String cf = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx";
        setStr(cf, "k1", "v1");
        setStr(cf, "k2", "v2");
        setStr(cf, "k3", "v3");
        setStr(cf, "k1", "v");
        System.out.println(getStr(cf, "k1"));
        Map<String, String> stringStringMap = multiGetStr(cf, Lists.newArrayList("k1", "k2", "k3"));
        Iterator<KV<String>> iter = iteratorGet(cf);
        iter.forEachRemaining((i) -> {
            System.out.println("iter "+i.toString());
        });
    }


    public static ColumnFamilyHandle getColumnFamilyHandle(String cf) {
        ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(cf);
        if (columnFamilyHandle == null) {
            synchronized (RocksDBExample.class) {
                try {
                    List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
                    List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<>();
                    // Only open the database once
                    if (db == null) {
                        List<byte[]> existingColumnFamilies = RocksDB.listColumnFamilies(getOptions(), cfdbPath);

                        for (byte[] cfName : existingColumnFamilies) {
                            columnFamilyDescriptors.add(new ColumnFamilyDescriptor(cfName, new ColumnFamilyOptions().optimizeLevelStyleCompaction()));
                        }

                        // 添加默认列族（如果不存在）
                        boolean defaultExists = false;
                        for (ColumnFamilyDescriptor cfd : columnFamilyDescriptors) {
                            if (Arrays.equals(cfd.getName(), RocksDB.DEFAULT_COLUMN_FAMILY)) {
                                defaultExists = true;
                                break;
                            }
                        }
                        if (!defaultExists) {
                            columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY,
                                    new ColumnFamilyOptions().optimizeLevelStyleCompaction()));
                        }

                        DBOptions dbOptions = new DBOptions()
                                .setCreateIfMissing(true)
                                .setCreateMissingColumnFamilies(true);
                        db = RocksDB.open(dbOptions, cfdbPath, columnFamilyDescriptors, columnFamilyHandleList);
                        for (int i = 0; i < columnFamilyDescriptors.size(); i++) {
                            columnFamilyHandles.put(new String(columnFamilyDescriptors.get(i).getName()), columnFamilyHandleList.get(i));
                        }
                    }


                    columnFamilyHandle = columnFamilyHandles.get(cf);
                    if (columnFamilyHandle == null) {
                        synchronized (RocksDBExample.class) {
                            if (columnFamilyHandle == null) {
                                // 数据库已打开，使用createColumnFamily动态创建新列族
                                ColumnFamilyOptions cfOpts = new ColumnFamilyOptions().optimizeLevelStyleCompaction();
                                ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(cf.getBytes(), cfOpts);
                                columnFamilyHandle = db.createColumnFamily(columnFamilyDescriptor);
                                columnFamilyHandles.put(cf, columnFamilyHandle);
                            }
                        }
                    }

                    // Update column family handles map
                    for (int i = 0; i < columnFamilyDescriptors.size(); i++) {
                        String cfName = new String(columnFamilyDescriptors.get(i).getName());
                        columnFamilyHandles.put(cfName, columnFamilyHandleList.get(i));
                    }

                    columnFamilyHandle = columnFamilyHandles.get(cf);
                } catch (RocksDBException e) {
                    throw new RuntimeException(e);
                }
            }
        }
        return columnFamilyHandle;
    }


    public static void setStr(String cf, String k, String v) {
        try {
            ColumnFamilyHandle cfHandle = getColumnFamilyHandle(cf);
            db.put(cfHandle, k.getBytes(StandardCharsets.UTF_8), v.getBytes(StandardCharsets.UTF_8));
        } catch (RocksDBException e) {
            throw new RuntimeException("Failed to put key-value pair", e);
        }
    }

    public static void delStr(String cf, String k) {
        try {
            ColumnFamilyHandle cfHandle = getColumnFamilyHandle(cf);
            db.delete(cfHandle, k.getBytes(StandardCharsets.UTF_8));
        } catch (RocksDBException e) {
            throw new RuntimeException("Failed to put key-value pair", e);
        }
    }

    public static void deleteColumnFamily(String cf) {
        if (cf == null || cf.isEmpty()) {
            throw new IllegalArgumentException("Column family name cannot be null or empty");
        }

        synchronized (RocksDBExample.class) {
            try {
                ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(cf);
                if (columnFamilyHandle != null) {
                    // 删除列族
                    db.dropColumnFamily(columnFamilyHandle);
                    // 关闭列族句柄
                    columnFamilyHandle.close();
                    // 从缓存中移除
                    columnFamilyHandles.remove(cf);
                    System.out.println("Successfully deleted column family: " + cf);
                } else {
                    System.out.println("Column family not found: " + cf);
                }
            } catch (RocksDBException e) {
                throw new RuntimeException("Failed to delete column family: " + cf, e);
            }
        }
    }

    public static String getStr(String cf, String k) {
        try {
            ColumnFamilyHandle cfHandle = getColumnFamilyHandle(cf);
            byte[] value = db.get(cfHandle, k.getBytes(StandardCharsets.UTF_8));
            return value != null ? new String(value, StandardCharsets.UTF_8) : null;
        } catch (RocksDBException e) {
            throw new RuntimeException("Failed to get value", e);
        }
    }

    public static Map<String, String> multiGetStr(String cf, List<String> ks) {
        Map<String, String> result = new HashMap<>();
        try {
            ColumnFamilyHandle cfHandle = getColumnFamilyHandle(cf);
            List<byte[]> keys = ks.stream()
                    .map(key -> key.getBytes(StandardCharsets.UTF_8))
                    .collect(Collectors.toList());
            List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<>();
            keys.forEach(k -> columnFamilyHandleList.add(cfHandle));
            List<byte[]> values = db.multiGetAsList(columnFamilyHandleList, keys);
            for (int i = 0; i < ks.size(); i++) {
                if (values.get(i) != null) {
                    result.put(ks.get(i), new String(values.get(i), StandardCharsets.UTF_8));
                }
            }
        } catch (RocksDBException e) {
            throw new RuntimeException("Failed to multi-get values", e);
        }
        return result;
    }

    @Data
    @AllArgsConstructor
    public final static class KV<T> {
        private String key;
        private T val;
    }


    public static Iterator<KV<String>> iteratorGet(String cf) {
        ColumnFamilyHandle cfHandle = getColumnFamilyHandle(cf);
        RocksIterator iter = db.newIterator(cfHandle);
        iter.seekToFirst();
        return new Iterator<>() {
            boolean first = true;

            @Override
            public boolean hasNext() {
                if (first) {
                    first = false;
                } else {
                    iter.next();
                }
                boolean valid = iter.isValid();
                if (!valid) {
                    iter.close();  // 确保资源释放
                }
                return valid;
            }

            @Override
            public KV<String> next() {

                if (!iter.isValid()) {  // 双重检查
                    throw new NoSuchElementException("----- Iterator has no more elements");
                }
                return new KV<>(new String(iter.key()), new String(iter.value()));
            }
        };
    }


    public static void testBasicOperate() {
        System.out.println("开始测试rocksdb的基本操作...");

        final Options options = getOptions();
        try (final RocksDB db = RocksDB.open(options, dbPath)) {
            List<byte[]> keys = Lists.newArrayList();
            keys.add("hello".getBytes());

            db.put("hello".getBytes(), "world".getBytes());
            byte[] value = db.get("hello".getBytes());
            System.out.format("Get('hello') = %s\n", new String(value));

            // write batch test
            try (final WriteOptions writeOpt = new WriteOptions()) {
                for (int i = 1; i <= 9; ++i) {
                    try (final WriteBatch batch = new WriteBatch()) {
                        for (int j = 1; j <= 9; ++j) {
                            batch.put(String.format("%dx%d", i, j).getBytes(),
                                    String.format("%d", i * j).getBytes());
                            keys.add(String.format("%dx%d", i, j).getBytes());
                        }
                        db.write(writeOpt, batch);
                    }
                }
            }

            System.out.println("multiGetAsList方法获取");
            List<byte[]> values = db.multiGetAsList(keys);
            for (int i = 0; i < keys.size(); i++) {
                System.out.println(String.format("key:%s,value:%s",
                        new String(keys.get(i)),
                        (values.get(i) != null ? new String(values.get(i)) : null)));
            }

            System.out.println("newIterator方法获取");
            RocksIterator iter = db.newIterator();
            for (iter.seekToFirst(); iter.isValid(); iter.next()) {
                System.out.println(String.format("key:%s,value:%s",
                        new String(iter.key()), new String(iter.value())));
            }

        } catch (RocksDBException e) {
            e.printStackTrace();
        }
    }

    // 使用特定的列族打开数据库，可以把列族理解为关系型数据库中的表(table)
    public static void testCustomColumnFamily() {
        System.out.println("测试自定义的列簇...");
        try (final ColumnFamilyOptions cfOpts = new ColumnFamilyOptions().optimizeLevelStyleCompaction()) {
            String cfName = "cf";
            // list of column family descriptors, first entry must always be default column family
            final List<ColumnFamilyDescriptor> cfDescriptors = Arrays.asList(
                    new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, cfOpts),
                    new ColumnFamilyDescriptor(cfName.getBytes(), cfOpts)
            );

            List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
            try (final DBOptions dbOptions = new DBOptions().setCreateIfMissing(true).setCreateMissingColumnFamilies(true);
                 final RocksDB db = RocksDB.open(dbOptions, cfdbPath, cfDescriptors, cfHandles)) {

                ColumnFamilyHandle cfHandle = cfHandles.stream().filter(x -> {
                    try {
                        return (new String(x.getName())).equals(cfName);
                    } catch (RocksDBException e) {
                        return false;
                    }
                }).collect(Collectors.toList()).get(0);

                try {
                    // put and get from non-default column family
                    db.put(cfHandles.get(1), new WriteOptions(), "key".getBytes(), "value".getBytes());

                    // atomic write
                    try (final WriteBatch wb = new WriteBatch()) {
                        wb.put(cfHandles.get(0), "key2".getBytes(),
                                "value2".getBytes());
                        wb.put(cfHandles.get(1), "key3".getBytes(),
                                "value3".getBytes());
                        wb.delete(cfHandles.get(1), "key".getBytes());
                        db.write(new WriteOptions(), wb);
                    }

                    System.out.println("newIterator方法获取");
                    //如果不传columnFamilyHandle，则获取默认的列簇，如果传了columnFamilyHandle，则获取指定列簇的
                    RocksIterator iter = db.newIterator(cfHandles.get(1));
                    for (iter.seekToFirst(); iter.isValid(); iter.next()) {
                        System.out.println(String.format("key:%s,value:%s",
                                new String(iter.key()), new String(iter.value())));
                    }

                    // drop column family
                    db.dropColumnFamily(cfHandles.get(1));

                } finally {
                    for (final ColumnFamilyHandle handle : cfHandles) {
                        handle.close();
                    }
                }
            }
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
    }

    public static void testTransaction() {
        System.out.println("测试事务开始...");
        try (final Options options = new Options()
                .setCreateIfMissing(true);
             final OptimisticTransactionDB txnDb =
                     OptimisticTransactionDB.open(options, txdbPath)) {

            try (final WriteOptions writeOptions = new WriteOptions();
                 final ReadOptions readOptions = new ReadOptions()) {

                System.out.println("=========================================");
                System.out.println("Demonstrates \"Read Committed\" isolation");
                readCommitted(txnDb, writeOptions, readOptions);
                iteratorReadData(txnDb);

                System.out.println("=========================================");
                System.out.println("Demonstrates \"Repeatable Read\" (Snapshot Isolation) isolation");
                repeatableRead(txnDb, writeOptions, readOptions);
                iteratorReadData(txnDb);

                System.out.println("=========================================");
                System.out.println("Demonstrates \"Read Committed\" (Monotonic Atomic Views) isolation");
                readCommitted_monotonicAtomicViews(txnDb, writeOptions, readOptions);
                iteratorReadData(txnDb);
            }
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
    }

    private static void iteratorReadData(RocksDB db) {
        System.out.println("newIterator方法获取");
        RocksIterator iter = db.newIterator();
        for (iter.seekToFirst(); iter.isValid(); iter.next()) {
            System.out.println(String.format("key:%s,value:%s",
                    new String(iter.key()), new String(iter.value())));
        }
    }

    /**
     * Demonstrates "Read Committed" isolation
     */
    private static void readCommitted(final OptimisticTransactionDB txnDb,
                                      final WriteOptions writeOptions, final ReadOptions readOptions)
            throws RocksDBException {
        final byte key1[] = "abc".getBytes(UTF_8);
        final byte value1[] = "def".getBytes(UTF_8);

        final byte key2[] = "xyz".getBytes(UTF_8);
        final byte value2[] = "zzz".getBytes(UTF_8);

        // Start a transaction
        try (final Transaction txn = txnDb.beginTransaction(writeOptions)) {
            // Read a key in this transaction
            byte[] value = txn.get(readOptions, key1);
            assert (value == null);

            // Write a key in this transaction
            txn.put(key1, value1);

            // Read a key OUTSIDE this transaction. Does not affect txn.
            value = txnDb.get(readOptions, key1);
            assert (value == null);

            // Write a key OUTSIDE of this transaction.
            // Does not affect txn since this is an unrelated key.
            // If we wrote key 'abc' here, the transaction would fail to commit.
            txnDb.put(writeOptions, key2, value2);

            // Commit transaction
            txn.commit();
        }
    }

    /**
     * Demonstrates "Repeatable Read" (Snapshot Isolation) isolation
     */
    private static void repeatableRead(final OptimisticTransactionDB txnDb,
                                       final WriteOptions writeOptions, final ReadOptions readOptions)
            throws RocksDBException {

        final byte key1[] = "ghi".getBytes(UTF_8);
        final byte value1[] = "jkl".getBytes(UTF_8);

        // Set a snapshot at start of transaction by setting setSnapshot(true)
        try (final OptimisticTransactionOptions txnOptions =
                     new OptimisticTransactionOptions().setSetSnapshot(true);
             final Transaction txn =
                     txnDb.beginTransaction(writeOptions, txnOptions)) {

            final Snapshot snapshot = txn.getSnapshot();

            // Write a key OUTSIDE of transaction
            txnDb.put(writeOptions, key1, value1);

            // Read a key using the snapshot.
            readOptions.setSnapshot(snapshot);
            final byte[] value = txn.getForUpdate(readOptions, key1, true);
            assert (value == null);

            try {
                // Attempt to commit transaction
                txn.commit();
                throw new IllegalStateException();
            } catch (final RocksDBException e) {
                // Transaction could not commit since the write outside of the txn
                // conflicted with the read!
                System.out.println(e.getStatus().getCode());
                assert (e.getStatus().getCode() == Status.Code.Busy);
            }

            txn.rollback();
        } finally {
            // Clear snapshot from read options since it is no longer valid
            readOptions.setSnapshot(null);
        }
    }

    /**
     * Demonstrates "Read Committed" (Monotonic Atomic Views) isolation
     * <p>
     * In this example, we set the snapshot multiple times.  This is probably
     * only necessary if you have very strict isolation requirements to
     * implement.
     */
    private static void readCommitted_monotonicAtomicViews(
            final OptimisticTransactionDB txnDb, final WriteOptions writeOptions,
            final ReadOptions readOptions) throws RocksDBException {

        final byte keyX[] = "x".getBytes(UTF_8);
        final byte valueX[] = "x".getBytes(UTF_8);

        final byte keyY[] = "y".getBytes(UTF_8);
        final byte valueY[] = "y".getBytes(UTF_8);

        try (final OptimisticTransactionOptions txnOptions =
                     new OptimisticTransactionOptions().setSetSnapshot(true);
             final Transaction txn =
                     txnDb.beginTransaction(writeOptions, txnOptions)) {

            // Do some reads and writes to key "x"
            Snapshot snapshot = txnDb.getSnapshot();
            readOptions.setSnapshot(snapshot);
            byte[] value = txn.get(readOptions, keyX);
            txn.put(keyX, valueX);

            // Do a write outside of the transaction to key "y"
            txnDb.put(writeOptions, keyY, valueY);

            // Set a new snapshot in the transaction
            txn.setSnapshot();
            snapshot = txnDb.getSnapshot();
            readOptions.setSnapshot(snapshot);

            // Do some reads and writes to key "y"
            // Since the snapshot was advanced, the write done outside of the
            // transaction does not conflict.
            value = txn.getForUpdate(readOptions, keyY, true);
            txn.put(keyY, valueY);

            // Commit.  Since the snapshot was advanced, the write done outside of the
            // transaction does not prevent this transaction from Committing.
            txn.commit();

        } finally {
            // Clear snapshot from read options since it is no longer valid
            readOptions.setSnapshot(null);
        }
    }


}
